"""
app.py
把训练好的 best.pth 封装成 /predict 接口
"""
import uvicorn, torch, io
from fastapi import FastAPI, File, UploadFile
from PIL import Image
from src.config import Config
from src.model import get_model
from src.dataset import HerbDataset   # 复用预处理

# 1. 加载模型
device = Config.device
model = get_model().to(device)
model.load_state_dict(torch.load("models/best.pth", map_location=device))
model.eval()
classes = ['猫', '狗', '其他']
transform = HerbDataset(Config.test_dir, Config.input_size, 'test').trans

# 2. 创建 FastAPI
app = FastAPI(title="HerbClassifier")

temperature = 2.0  # 调整温度参数

# app.py  仅修改返回值
@app.post("/predict")
async def predict(file: UploadFile = File(...)):
    try:
        img = Image.open(io.BytesIO(await file.read())).convert('RGB')
        x = transform(img).unsqueeze(0).to(device)
        with torch.no_grad():
            # prob, idx = torch.max(torch.softmax(model(x), dim=1), 1)
            logits = model(x)
            # 应用temperature scaling
            scaled_logits = logits / temperature
            probs = torch.softmax(scaled_logits, dim=1)
            prob, idx = torch.max(probs, 1)

            # 更严格的置信度阈值
            if prob.item() < 0.8:
                return {"class": "unknown", "confidence": prob.item()}
        return {"class": classes[idx], "confidence": prob.item()}
    except Exception as e:
        return {"error": str(e)}


@app.post("/predictAll")
async def predict(file: UploadFile = File(...)):
    try:
        img = Image.open(io.BytesIO(await file.read())).convert('RGB')
        x = transform(img).unsqueeze(0).to(device)
        with torch.no_grad():
            probs = torch.softmax(model(x), dim=1).cpu().tolist()[0]
        return {"class": classes[probs.index(max(probs))],
                "probs": dict(zip(classes, probs))}
    except Exception as e:
        return {"error": str(e)}

# 3. 启动
if __name__ == "__main__":
    uvicorn.run("app:app", host="0.0.0.0", port=8410)